import serial
import minimalmodbus
import time
import threading
import sys
from flask import Flask, request, jsonify, send_file
from flask_cors import CORS
import io
import csv
from datetime import datetime

app = Flask(__name__)
# 优化跨域配置
CORS(app, resources={
    r"/api/*": {
        "origins": "*",
        "methods": ["GET", "POST", "OPTIONS"],
        "allow_headers": ["Content-Type"]
    }
})

# 全局状态变量
system_state = "待机"
measurement_data = []
is_measuring = False
collector = None
motor = None
sensor = None
stepper = None
shared_serial = None
serial_lock = threading.Lock()

# 设备配置参数
config = {
    'port': 'COM6',
    'baudrate': 115200,
    'motor_id': 1,
    'sensor_address': 2,
    'stepper_id': 3,
    'stepper_cycles': 2,
    'sample_interval': 0.05
}

# 处理预检请求
@app.before_request
def handle_options():
    if request.method == 'OPTIONS':
        response = app.make_default_options_response()
        response.headers['Access-Control-Allow-Headers'] = 'Content-Type'
        return response

class LingkongMotor:
    def __init__(self, serial_port, motor_id=0x01, lock=None):
        if serial_port is None:
            raise ValueError("必须提供共享 serial_port 实例")
        self.ser = serial_port
        self.motor_id = motor_id
        self.low_speed = 20
        self.zero_angle_target = 0
        self.lock = lock or threading.Lock()

    def _calc_checksum(self, data):
        return sum(data) & 0xFF

    def _send_command(self, cmd, data_bytes=b'', read_len=32, post_delay=0.002):
        frame_cmd = [0x3E, cmd, self.motor_id, len(data_bytes)]
        frame_cmd_checksum = self._calc_checksum(frame_cmd)
        frame_cmd.append(frame_cmd_checksum)

        full_frame = bytes(frame_cmd) + data_bytes
        if data_bytes:
            data_checksum = self._calc_checksum(data_bytes)
            full_frame += bytes([data_checksum])

        with self.lock:
            try:
                try:
                    self.ser.reset_input_buffer()
                    self.ser.reset_output_buffer()
                except Exception:
                    pass

                self.ser.write(full_frame)
                self.ser.flush()
                time.sleep(post_delay)
                resp = self.ser.read(read_len)
                return resp
            except Exception as e:
                print(f"[Motor] 发送命令失败: {e}")
                raise

    def get_current_angle(self):
        try:
            response = self._send_command(0x92, b'', read_len=32)
        except Exception as e:
            print(f"[Motor] 获取角度命令出错: {e}")
            return None

        if response and len(response) >= 14:
            try:
                angle_bytes = response[5:13]
                current_angle_raw = int.from_bytes(angle_bytes, byteorder='little', signed=True)
                return current_angle_raw * 0.01
            except Exception as e:
                print(f"[Motor] 解析角度失败: {e}")
                return None
        return None

    def move_to_zero(self, timeout=300):
        print("开始回归零位...")
        target_bytes = int(self.zero_angle_target).to_bytes(8, byteorder='little', signed=True)
        speed_value = int(self.low_speed * 100)
        speed_bytes = int(speed_value).to_bytes(4, byteorder='little', signed=False)

        try:
            self._send_command(0xA4, target_bytes + speed_bytes)
        except Exception as e:
            print(f"[Motor] 发送回零命令失败: {e}")
            return False

        start_time = time.time()
        while time.time() - start_time < timeout:
            current = self.get_current_angle()
            if current is not None and abs(current) <= 0.1:
                print("已到达零位")
                return True
            time.sleep(0.2)
        print("回零超时")
        return False

    def rotate_360(self, timeout=300):
        print("开始旋转360度...")
        target = 36000  # 360.00度（原始值）
        target_bytes = int(target).to_bytes(8, byteorder='little', signed=True)
        speed_value = int(self.low_speed * 10)  # 速度原始值
        speed_bytes = int(speed_value).to_bytes(4, byteorder='little', signed=False)

        try:
            self._send_command(0xA4, target_bytes + speed_bytes)
        except Exception as e:
            print(f"[Motor] 发送旋转命令失败: {e}")
            return False

        start_time = time.time()
        while time.time() - start_time < timeout:
            current = self.get_current_angle()
            if current is not None:
                ang = current % 360
                if abs(ang - 0) <= 0.1:
                    print("已完成360度旋转（归一化为0）")
                    return True
            time.sleep(0.2)
        print("旋转超时")
        return False

    def close(self):
        pass


class LE100485Sensor:
    def __init__(self, serial_port, sensor_address=2, lock=None):
        if serial_port is None:
            raise ValueError("必须提供共享 serial_port 实例")
        self.serial_port = serial_port
        self.lock = lock or threading.Lock()
        self.sensor_address = sensor_address

        self.instrument = minimalmodbus.Instrument(serial_port, slaveaddress=self.sensor_address)
        self.instrument.mode = minimalmodbus.MODE_RTU
        self.instrument.serial = self.serial_port
        self.instrument.close_port_after_each_call = False
        self.instrument.clear_buffers_before_each_transaction = True
        self.instrument.debug = False

        self.distance_register = 0x003B  # 距离寄存器地址
        self.resolution = 0.001  # 分辨率：1mm = 1000单位

    def read_distance(self, register_address=None):
        if register_address is None:
            register_address = self.distance_register
        with self.lock:
            try:
                raw_data = self.instrument.read_register(
                    registeraddress=register_address,
                    functioncode=3,
                    signed=False
                )
                return raw_data * self.resolution
            except Exception as e:
                print(f"[Sensor] 读取失败: {e}")
                return None

    def close(self):
        pass


class StepperMotor:
    """步进电机控制类（地址3）"""
    def __init__(self, serial_port, motor_id=3, lock=None, 细分=10000):
        self.ser = serial_port
        self.motor_id = motor_id  # 步进电机地址3
        self.细分 = 细分  # 细分参数，默认为10000
        self.lock = lock or threading.Lock()

    def _send_modbus_command(self, register, value):
        """发送Modbus RTU写单个寄存器命令（功能码06）"""
        with self.lock:
            try:
                # 构建命令帧：地址 + 功能码06 + 寄存器高8位 + 寄存器低8位 + 数据高8位 + 数据低8位 + 校验和
                cmd = bytearray()
                cmd.append(self.motor_id)  # 设备地址
                cmd.append(0x06)  # 功能码06（写单个寄存器）
                
                # 寄存器地址（2字节）
                cmd.extend(register.to_bytes(2, byteorder='big'))
                
                # 写入值（2字节）
                cmd.extend(value.to_bytes(2, byteorder='big'))
                
                # 计算CRC校验和
                crc = self._calc_crc(cmd)
                cmd.extend(crc)
                
                # 发送命令
                self.ser.reset_input_buffer()
                self.ser.reset_output_buffer()
                self.ser.write(cmd)
                self.ser.flush()
                
                # 等待响应（8字节）
                time.sleep(0.01)
                resp = self.ser.read(8)
                return len(resp) == 8  # 响应正确返回True
            except Exception as e:
                print(f"[Stepper] 发送命令失败: {e}")
                return False

    def _calc_crc(self, data):
        """计算Modbus CRC16校验和"""
        crc = 0xFFFF
        for byte in data:
            crc ^= byte
            for _ in range(8):
                if crc & 0x0001:
                    crc >>= 1
                    crc ^= 0xA001
                else:
                    crc >>= 1
        return crc.to_bytes(2, byteorder='little')

    def rotate_cycles(self, cycles, speed=300, accel=50, decel=50, timeout=300):
        """
        控制步进电机旋转指定圈数
        cycles: 旋转圈数（正数为正转，负数为反转）
        """
        print(f"开始控制步进电机旋转 {cycles} 圈...")
        
        # 计算目标位置（1圈 = 细分值）
        target_position = int(cycles * self.细分)
        
        # 设置PR0为绝对位置模式（0001）
        if not self._send_modbus_command(0x6200, 0x0001):
            print("设置PR0模式失败")
            return False
        
        # 处理目标位置（分高低位）
        position_high = (target_position >> 16) & 0xFFFF
        position_low = target_position & 0xFFFF
        if not self._send_modbus_command(0x6201, position_high):
            print("设置PR0位置高位失败")
            return False
        if not self._send_modbus_command(0x6202, position_low):
            print("设置PR0位置低位失败")
            return False
        
        # 设置速度、加速度、减速度
        if not self._send_modbus_command(0x6203, speed):
            print("设置PR0速度失败")
            return False
        if not self._send_modbus_command(0x6204, accel):
            print("设置PR0加速度失败")
            return False
        if not self._send_modbus_command(0x6205, decel):
            print("设置PR0减速度失败")
            return False
        
        # 触发PR0运行
        if not self._send_modbus_command(0x6002, 0x0010):
            print("触发PR0运行失败")
            return False
        
        # 等待完成
        start_time = time.time()
        while time.time() - start_time < timeout:
            time.sleep(1)
            estimated_time = abs(cycles * self.细分) / (speed * 1.0)
            if time.time() - start_time > estimated_time + 2:
                print(f"步进电机完成 {cycles} 圈旋转")
                return True
        
        print("步进电机旋转超时")
        return False

    def emergency_stop(self):
        """急停命令"""
        return self._send_modbus_command(0x6002, 0x0040)


class DataCollector:
    def __init__(self, motor, sensor):
        self.motor = motor
        self.sensor = sensor
        self.data = []
        self.running = False
        self.thread = None

    def _collect_loop(self, interval):
        global measurement_data
        while self.running:
            angle = self.motor.get_current_angle()
            distance = self.sensor.read_distance()

            if angle is not None and distance is not None:
                normalized_angle = angle % 360
                self.data.append((normalized_angle, distance))
                measurement_data.append((normalized_angle, distance))
                print(f"角度: {normalized_angle:.2f}°  距离: {distance:.5f}mm", end='\r', flush=True)

            time.sleep(interval)

    def start(self, interval=0.05):
        if self.running:
            return
        self.data = []
        self.running = True
        self.thread = threading.Thread(target=self._collect_loop, args=(interval,), daemon=True)
        self.thread.start()

    def stop(self):
        self.running = False
        if self.thread:
            self.thread.join(timeout=2)

    def save_data(self, filename="measurement_data.csv"):
        with open(filename, 'w', encoding='utf-8') as f:
            f.write("角度(度),距离(mm)\n")
            for angle, distance in self.data:
                f.write(f"{angle:.4f},{distance:.5f}\n")
        print(f"\n数据已保存到 {filename}，共 {len(self.data)} 条记录")


# API接口实现
@app.route('/api/connect', methods=['POST'])
def connect():
    global shared_serial, motor, sensor, stepper, system_state
    
    try:
        # 获取前端配置
        data = request.json
        if data:
            config.update(data)
            
        # 初始化共享串口
        shared_serial = serial.Serial(
            port=config['port'],
            baudrate=config['baudrate'],
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            bytesize=serial.EIGHTBITS,
            timeout=1,
            write_timeout=5
        )

        # 配置RS485模式
        try:
            if hasattr(serial, "rs485"):
                rs485_settings = serial.rs485.RS485Settings(
                    rts_level_for_tx=True,
                    rts_level_for_rx=False,
                    delay_before_tx=0.002,
                    delay_before_rx=0.002
                )
                shared_serial.rs485_mode = rs485_settings
                print("[Main] 已启用 serial.rs485_mode")
        except Exception as e:
            print(f"[Main] 无法启用 rs485_mode（忽略）: {e}")

        # 初始化设备
        motor = LingkongMotor(serial_port=shared_serial, motor_id=config['motor_id'], lock=serial_lock)
        sensor = LE100485Sensor(serial_port=shared_serial, sensor_address=config['sensor_address'], lock=serial_lock)
        stepper = StepperMotor(serial_port=shared_serial, motor_id=config['stepper_id'], lock=serial_lock)
        
        system_state = "已连接"
        return jsonify({"status": "success", "message": "设备连接成功"})
    except Exception as e:
        import traceback
        traceback.print_exc()
        return jsonify({"status": "error", "message": f"连接失败: {str(e)}"})


@app.route('/api/start', methods=['POST'])
def start_measure():
    global is_measuring, collector, system_state, measurement_data, stepper, motor, sensor
    
    try:
        if is_measuring:
            return jsonify({"status": "error", "message": "正在测量中"})
        
        # 检查设备是否已初始化
        if not stepper or not motor or not sensor:
            return jsonify({"status": "error", "message": "设备未连接，请先完成连接"})
            
        # 重置测量数据
        measurement_data = []
        
        # 初始化数据采集器
        collector = DataCollector(motor, sensor)
        
        # 第一步：控制步进电机旋转指定圈数
        system_state = "步进电机旋转中"
        if not stepper.rotate_cycles(config['stepper_cycles']):
            system_state = "错误"
            return jsonify({"status": "error", "message": "步进电机旋转失败"})
        
        # 第二步：关节电机回零
        system_state = "关节电机回零中"
        if not motor.move_to_zero():
            system_state = "错误"
            return jsonify({"status": "error", "message": "关节电机无法回归零位"})
        
        # 第三步：启动数据采集并控制关节电机旋转
        system_state = "测量中"
        collector.start(interval=config['sample_interval'])
        is_measuring = True
        
        # 在后台线程中执行旋转和采集
        def measurement_task():
            global is_measuring, system_state
            if not motor.rotate_360():
                system_state = "错误"
                collector.stop()
                is_measuring = False
        
        threading.Thread(target=measurement_task, daemon=True).start()
        
        return jsonify({"status": "success", "message": "开始测量"})
    except Exception as e:
        import traceback
        traceback.print_exc()
        return jsonify({"status": "error", "message": f"启动失败: {str(e)}"}), 500


@app.route('/api/stop', methods=['POST'])
def stop_measure():
    global is_measuring, system_state, collector
    
    if collector:
        collector.stop()
    is_measuring = False
    system_state = "已停止"
    return jsonify({"status": "success", "message": "已停止测量"})


@app.route('/api/emergency', methods=['POST'])
def emergency_stop():
    global is_measuring, system_state, collector
    
    if collector:
        collector.stop()
    if stepper:
        stepper.emergency_stop()
    is_measuring = False
    system_state = "急停"
    return jsonify({"status": "success", "message": "已执行急停"})


@app.route('/api/reset', methods=['POST'])
def reset_system():
    global is_measuring, system_state, collector, measurement_data
    
    if collector:
        collector.stop()
    measurement_data = []
    is_measuring = False
    
    # 尝试回零
    if motor:
        system_state = "复位中"
        if motor.move_to_zero():
            system_state = "待机"
        else:
            system_state = "复位失败"
    else:
        system_state = "待机"
        
    return jsonify({"status": "success", "message": "系统已复位"})


@app.route('/api/status', methods=['GET'])
def get_status():
    global measurement_data, system_state, motor, sensor
    
    angle = None
    distance = None
    
    if motor:
        angle = motor.get_current_angle()
    if sensor:
        distance = sensor.read_distance()
    
    return jsonify({
        "system": system_state,
        "dataCount": len(measurement_data) if collector else 0,
        "distance": distance,
        "angle": angle
    })


@app.route('/api/data', methods=['GET'])
def get_data():
    global measurement_data
    return jsonify({
        "angles": [d[0] for d in measurement_data],
        "distances": [d[1] for d in measurement_data]
    })


@app.route('/api/export', methods=['POST'])
def export_data():
    global measurement_data
    
    if not measurement_data:
        return jsonify({"status": "error", "message": "没有数据可导出"}), 400
    
    # 创建CSV文件
    output = io.StringIO()
    writer = csv.writer(output)
    writer.writerow(["角度(度)", "距离(mm)"])
    for angle, distance in measurement_data:
        writer.writerow([f"{angle:.4f}", f"{distance:.5f}"])
    
    # 准备下载
    output.seek(0)
    filename = f"measurement_data_{datetime.now().strftime('%Y%m%d_%H%M%S')}.csv"
    
    return send_file(
        io.BytesIO(output.getvalue().encode('utf-8')),
        mimetype="text/csv",
        as_attachment=True,
        download_name=filename
    )


if __name__ == "__main__":
    # 启动Flask服务器
    app.run(host='0.0.0.0', port=5000, debug=True)